// This code is licensed under New BSD Licence. For details see project page at
// http://code.google.com/p/hrissan-cocoa-helpers/source/checkout

#ifndef __RFC_RC4_H__
#define __RFC_RC4_H__

namespace rfc
{

class rc4
{
	unsigned char state[256];
	unsigned char x;
	unsigned char y;

	void prepare_key(const unsigned char * key_data_ptr, int key_data_len)
	{
		{
			for(int counter = 0; counter < 256; counter++)
				state[counter] = counter;
		}
		x = 0;
		y = 0;
		unsigned char index1 = 0;
		unsigned char index2 = 0;
		{
			for(int counter = 0; counter < 256; counter++)
			{
				index2 = (key_data_ptr[index1] + state[counter] + index2) % 256;
				std::swap(state[counter], state[index2]);

				index1 = (index1 + 1) % key_data_len;
			}
		}
	}
public:
	rc4(const unsigned char * key_data_ptr, int key_data_len)
	{
		prepare_key(key_data_ptr, key_data_len);
	}
	const unsigned char * get_state()const
	{
		return state;
	}
	void set_state( unsigned char* new_state, int new_x, int new_y )
	{
		for ( int i = 0; i <= 255; i++ )
			state[i] = new_state[i];
		
		x = new_x;
		y = new_y;
	}
	int get_x()const
	{
		return x;
	}
	int get_y()const
	{
		return x;
	}
	void encrypt(unsigned char * buffer_ptr, int buffer_len)
	{
		for(int counter = 0; counter < buffer_len; counter ++)
		{
			x = (x + 1) % 256;
			y = (state[x] + y) % 256;
			std::swap(state[x], state[y]);

			unsigned char xorIndex = (state[x] + state[y]) % 256;

			if( buffer_ptr )
				buffer_ptr[counter] ^= state[xorIndex];
		}
	}

	void skip(int octets)
	{
		for(int counter = 0; counter < octets*4; counter ++)
		{
			x = (x + 1) % 256;
			y = (state[x] + y) % 256;
			std::swap(state[x], state[y]);
		}
	}


	void encrypt(unsigned char * buffer_ptr, int buffer_len, int x)
	{
		for(int counter = 0; counter < buffer_len; counter ++)
		{
			x = (x + 1) % 256;
			y = (state[x] + y) % 256;
			std::swap(state[x], state[y]);

			unsigned char xorIndex = (state[x] + state[y]) % 256;

			if( buffer_ptr )
				buffer_ptr[counter] ^= state[xorIndex];
		}
	}
};

} //namespace rfc

#endif // __RFC_RC4_H__
